# encoding: utf-8

import numpy
import pyomo.environ as pyomo

class q_adjust_model(object):
    def __init__(self, delta_q_range, delta_q):
        #
        self.delta_q_range = delta_q_range

        self.n_devices = len(self.delta_q_range)

        self.delta_q = delta_q

    def q_assign_objective(self, model):
        q_adjust = [0] * len(self.delta_q_range)
        for i in range(0, len(self.delta_q_range)):
            q_adjust[i] = model.x[i].value

        adjust_cost = 0.0
        for i in range(0, len(self.delta_q_range)):
            adjust_cost = adjust_cost + 1/self.delta_q_range[i]*q_adjust[i]

        return adjust_cost

    def q_assign_rule1(self, model):
        return sum(model.x[j] * 1 for j in model.A) == self.delta_q

    def q_assign_rule2(self, model, i):
        return model.x[i] <= self.delta_q_range[i]

    def get_model(self):

        model = pyomo.ConcreteModel()

        data = {}
        # Initial data, lower and upper boundary
        for i in range(0, self.n_devices):
            data[i] = 0.0  #

        upper_bound = 1000#max(self.delta_q_range)
        lower_bound = 0

        model.s = pyomo.Set(initialize=range(0, self.n_devices))
        model.x = pyomo.Var(model.s, bounds=(lower_bound, upper_bound), initialize=data)

        model.obj = pyomo.Objective(rule=self.q_assign_objective, sense=pyomo.minimize)

        model.A = pyomo.RangeSet(0, self.n_devices-1)

        model.capacity1 = pyomo.Constraint(model.A, rule=self.q_assign_rule1)

        model.capacity2 = pyomo.Constraint(model.A, rule=self.q_assign_rule2)

        return model

